import sqlite3
from typing import List, Dict, Any, Optional, Union


class SQLiteDB:
    """
    SQLite 数据库操作工具类

    提供连接、创建表、增删改查等常用操作
    """

    def __init__(self, db_path: str = ":memory:"):
        """
        初始化数据库连接

        :param db_path: 数据库文件路径，默认为内存数据库
        """
        self.db_path = db_path
        self.connection = None
        self.cursor = None

    def connect(self) -> None:
        """连接到数据库"""
        try:
            self.connection = sqlite3.connect(self.db_path)
            self.cursor = self.connection.cursor()
            print(f"成功连接到数据库: {self.db_path}")
        except sqlite3.Error as e:
            print(f"连接数据库失败: {e}")
            raise

    def close(self) -> None:
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()
            print("数据库连接已关闭")

    def execute(self, sql: str, params: Union[tuple, Dict[str, Any], None] = None) -> None:
        """
        执行SQL语句

        :param sql: SQL语句
        :param params: 参数，可以是元组或字典
        """
        try:
            if params:
                self.cursor.execute(sql, params)
            else:
                self.cursor.execute(sql)
            self.connection.commit()
        except sqlite3.Error as e:
            self.connection.rollback()
            print(f"执行SQL失败: {e}\nSQL: {sql}")
            raise

    def executemany(self, sql: str, params_list: List[Union[tuple, Dict[str, Any]]]) -> None:
        """
        批量执行SQL语句

        :param sql: SQL语句
        :param params_list: 参数列表，每个元素可以是元组或字典
        """
        try:
            self.cursor.executemany(sql, params_list)
            self.connection.commit()
        except sqlite3.Error as e:
            self.connection.rollback()
            print(f"批量执行SQL失败: {e}\nSQL: {sql}")
            raise

    def create_table(self, table_name: str, columns: Dict[str, str]) -> None:
        """
        创建表

        :param table_name: 表名
        :param columns: 列定义字典，格式为 {"列名": "数据类型 约束"}
        """
        columns_sql = ", ".join([f"{name} {definition}" for name, definition in columns.items()])
        sql = f"CREATE TABLE IF NOT EXISTS {table_name} ({columns_sql})"
        self.execute(sql)
        print(f"表 {table_name} 创建成功")

    def insert(self, table_name: str, data: Dict[str, Any]) -> int:
        """
        插入单条数据

        :param table_name: 表名
        :param data: 数据字典，键为列名，值为数据
        :return: 插入行的rowid
        """
        columns = ", ".join(data.keys())
        placeholders = ", ".join(["?"] * len(data))
        sql = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"
        self.execute(sql, tuple(data.values()))
        return self.cursor.lastrowid

    def insert_many(self, table_name: str, data_list: List[Dict[str, Any]]) -> None:
        """
        批量插入数据

        :param table_name: 表名
        :param data_list: 数据字典列表
        """
        if not data_list:
            return

        columns = ", ".join(data_list[0].keys())
        placeholders = ", ".join(["?"] * len(data_list[0]))
        sql = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"

        params = [tuple(data.values()) for data in data_list]
        self.executemany(sql, params)

    def update(self, table_name: str, data: Dict[str, Any], condition: str, params: tuple = ()) -> int:
        """
        更新数据

        :param table_name: 表名
        :param data: 要更新的数据字典
        :param condition: WHERE条件语句
        :param params: WHERE条件的参数
        :return: 影响的行数
        """
        set_clause = ", ".join([f"{key} = ?" for key in data.keys()])
        sql = f"UPDATE {table_name} SET {set_clause} WHERE {condition}"
        self.execute(sql, tuple(data.values()) + params)
        return self.cursor.rowcount

    def delete(self, table_name: str, condition: str, params: tuple = ()) -> int:
        """
        删除数据

        :param table_name: 表名
        :param condition: WHERE条件语句
        :param params: WHERE条件的参数
        :return: 影响的行数
        """
        sql = f"DELETE FROM {table_name} WHERE {condition}"
        self.execute(sql, params)
        return self.cursor.rowcount

    def fetch_one(self, table_name: str, columns: List[str] = None,
                  condition: str = None, params: tuple = ()) -> Optional[Dict[str, Any]]:
        """
        查询单条数据

        :param table_name: 表名
        :param columns: 要查询的列名列表，None表示所有列
        :param condition: WHERE条件语句
        :param params: WHERE条件的参数
        :return: 单行数据字典，没有结果返回None
        """
        cols = "*" if columns is None else ", ".join(columns)
        sql = f"SELECT {cols} FROM {table_name}"
        if condition:
            sql += f" WHERE {condition}"
        sql += " LIMIT 1"

        self.execute(sql, params)
        row = self.cursor.fetchone()

        if row is None:
            return None

        if columns is None:
            columns = [desc[0] for desc in self.cursor.description]

        return dict(zip(columns, row))

    def fetch_all(self, table_name: str, columns: List[str] = None,
                  condition: str = None, params: tuple = ()) -> List[Dict[str, Any]]:
        """
        查询多条数据

        :param table_name: 表名
        :param columns: 要查询的列名列表，None表示所有列
        :param condition: WHERE条件语句
        :param params: WHERE条件的参数
        :return: 多行数据字典列表
        """
        cols = "*" if columns is None else ", ".join(columns)
        sql = f"SELECT {cols} FROM {table_name}"
        if condition:
            sql += f" WHERE {condition}"

        self.execute(sql, params)
        rows = self.cursor.fetchall()

        if not rows:
            return []

        if columns is None:
            columns = [desc[0] for desc in self.cursor.description]

        return [dict(zip(columns, row)) for row in rows]

    def count(self, table_name: str, condition: str = None, params: tuple = ()) -> int:
        """
        统计行数

        :param table_name: 表名
        :param condition: WHERE条件语句
        :param params: WHERE条件的参数
        :return: 行数
        """
        sql = f"SELECT COUNT(*) FROM {table_name}"
        if condition:
            sql += f" WHERE {condition}"

        self.execute(sql, params)
        return self.cursor.fetchone()[0]

    def table_exists(self, table_name: str) -> bool:
        """
        检查表是否存在

        :param table_name: 表名
        :return: 表是否存在
        """
        sql = "SELECT name FROM sqlite_master WHERE type='table' AND name=?"
        self.execute(sql, (table_name,))
        return self.cursor.fetchone() is not None

    def __enter__(self):
        """支持上下文管理"""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """支持上下文管理"""
        self.close()


# 使用示例
if __name__ == "__main__":
    # 创建内存数据库
    with SQLiteDB(":memory:") as db:
        # 创建表
        db.create_table("users", {
            "id": "INTEGER PRIMARY KEY AUTOINCREMENT",
            "name": "TEXT NOT NULL",
            "age": "INTEGER",
            "email": "TEXT UNIQUE"
        })

        # 插入数据
        user1 = {"name": "Alice", "age": 25, "email": "alice@example.com"}
        user2 = {"name": "Bob", "age": 30, "email": "bob@example.com"}
        db.insert("users", user1)
        db.insert_many("users", [user2])

        # 查询数据
        all_users = db.fetch_all("users")
        print("所有用户:", all_users)

        # 条件查询
        alice = db.fetch_one("users", condition="name = ?", params=("Alice",))
        print("Alice的信息:", alice)

        # 更新数据
        db.update("users", {"age": 26}, "name = ?", ("Alice",))

        # 删除数据
        db.delete("users", "name = ?", ("Bob",))

        # 再次查询
        remaining_users = db.fetch_all("users")
        print("剩余用户:", remaining_users)

        # 统计
        user_count = db.count("users")
        print(f"用户总数: {user_count}")
